''':cvar
1.使用cifar2,处理数据
2.将cifar2数据读入管道
3.使用ResNet模型处理数据
4.自定义参数,将模型调优
5.绘制训练和测试集的准确率曲线和代价函数曲线
'''

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings
warnings.filterwarnings('ignore')
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras import losses,models,layers,Sequential,optimizers
from tensorflow.keras.layers import Dense,Activation,BatchNormalization,GlobalAveragePooling2D,MaxPooling2D,Conv2D
import matplotlib.pyplot as plt

class CellBlock(layers.Layer):
    def __init__(self,filter_num,strides = 1):
        super(CellBlock, self).__init__()

        self.conv1 = Conv2D(filter_num,(3,3),strides=strides,padding='same')
        self.bn1 = BatchNormalization()
        self.relu = Activation('relu')

        self.conv2 = Conv2D(filter_num,(3,3),strides=1,padding='same')
        self.bn2 = BatchNormalization()

        if strides!=1:
            self.residual = Conv2D(filter_num,(1,1),strides=strides)
        else:
            self.residual = lambda x:x

    def call(self, inputs, **kwargs):

        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.relu(x)

        x = self.conv2(x)
        x = self.bn2(x)

        r = self.residual(inputs)

        x = layers.add([x,r])
        x = tf.nn.relu(x)

        return x

class ResNet(models.Model):
    def __init__(self,layers_dims,nb_classes):
        super(ResNet, self).__init__()

        self.stem = Sequential([
            Conv2D(64,(7,7),strides=(2,2),padding='same'),
            BatchNormalization(),
            Activation('relu'),
            MaxPooling2D((3,3),strides=(2,2),padding='same')
        ])

        self.layer1 = self.build_cellblock(64,layers_dims[0])
        self.layer2 = self.build_cellblock(128,layers_dims[1],strides=2)
        self.layer3 = self.build_cellblock(256,layers_dims[2],strides=2)
        self.layer4 = self.build_cellblock(512,layers_dims[3],strides=2)

        self.avgpool = GlobalAveragePooling2D()
        self.fc = Dense(nb_classes,activation='sigmoid')

    def call(self, inputs, training=None, mask=None):

        x = self.stem(inputs)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = self.fc(x)

        return x

    def build_cellblock(self,filter_num,blocks,strides=1):
        res_cellblock= Sequential()
        res_cellblock.add(CellBlock(filter_num,strides=strides))

        for _ in range(1,blocks):
            res_cellblock.add(CellBlock(filter_num,strides=1))

        return res_cellblock

def build_resnet(resname,nb_classes):
    res_config = {'resnet18':[2,2,2,2],
                  'resnet34':[3,3,6,3]}
    return ResNet(res_config[resname],nb_classes)


def load_imag(path):
    label = tf.constant(1,tf.int32)if tf.strings.regex_full_match(path,'.*automobile.*')else tf.constant(0,tf.int32)

    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img)
    img = tf.image.resize(img,size=(32,32))/255.0

    return img,label

if __name__ == '__main__':

    train = tf.data.Dataset.list_files('../data/cifar2/train/*/*.jpg')\
            .map(load_imag,tf.data.experimental.AUTOTUNE).shuffle(buffer_size=1000).batch(100)\
            .prefetch(tf.data.experimental.AUTOTUNE)

    test = tf.data.Dataset.list_files('../data/cifar2/test/*/*.jpg')\
            .map(load_imag,tf.data.experimental.AUTOTUNE).batch(100)\
            .prefetch(tf.data.experimental.AUTOTUNE)

    model = build_resnet('resnet18',nb_classes=1)

    model.build(input_shape=(None,32,32,3))

    model.summary()

    model.compile(loss = losses.binary_crossentropy,
                  optimizer=optimizers.Adam(),
                  metrics=['accuracy'])

    history = model.fit(train,epochs=6,validation_data=test,workers=4)

    train_loss = history.history['loss']
    test_loss = history.history['val_loss']

    train_acc = history.history['accuracy']
    test_acc = history.history['val_accuracy']

    plt.figure(figsize=(10,5))

    plt.subplot(121)
    plt.title('Train_loss And Test_loss')
    plt.plot(train_loss,label='train_loss')
    plt.plot(test_loss,label='test_loss')
    plt.legend()

    plt.subplot(122)
    plt.title('Train_acc And Test_acc')
    plt.plot(train_acc,label='train_acc')
    plt.plot(test_acc,label='test_acc')
    plt.legend()
    plt.show()